"""
基于Neo4j的代码推荐系统
利用调用关系数据提供智能代码补全建议
"""

import os
import json
from typing import List, Dict, Optional, Tuple
from neo4j import GraphDatabase, Driver
from dotenv import load_dotenv


class CodeRecommender:
    """基于Neo4j调用关系的代码推荐器"""
    
    def __init__(self, driver: Driver):
        self.driver = driver
    
    def get_function_recommendations(self, 
                                   current_function: str, 
                                   context_file: str = None,
                                   limit: int = 10) -> List[Dict]:
        """
        基于当前函数获取推荐调用的函数
        
        Args:
            current_function: 当前函数名或限定名
            context_file: 当前文件路径（可选）
            limit: 返回推荐数量限制
            
        Returns:
            推荐函数列表，包含函数信息和推荐理由
        """
        with self.driver.session() as session:
            # 查询1: 基于直接调用关系的推荐
            direct_calls_query = """
            MATCH (current:Function)-[c:CALLS]->(recommended:Function)
            WHERE current.name CONTAINS $function_name 
               OR current.id CONTAINS $function_name
            WITH recommended, count(c) as call_frequency
            ORDER BY call_frequency DESC
            LIMIT $limit
            RETURN recommended.id as function_id,
                   recommended.name as function_name,
                   recommended.file as function_file,
                   call_frequency,
                   'direct_call' as recommendation_type,
                   '经常被当前函数调用' as reason
            """
            
            # 查询2: 基于相似函数的推荐（调用相同函数的函数）
            similar_functions_query = """
            MATCH (current:Function)-[:CALLS]->(common:Function)<-[:CALLS]-(similar:Function)
            WHERE current.name CONTAINS $function_name 
               OR current.id CONTAINS $function_name
               AND similar <> current
            WITH similar, count(common) as common_calls
            ORDER BY common_calls DESC
            LIMIT $limit
            RETURN similar.id as function_id,
                   similar.name as function_name,
                   similar.file as function_file,
                   common_calls as call_frequency,
                   'similar_function' as recommendation_type,
                   '与当前函数调用相似函数' as reason
            """
            
            # 查询3: 基于文件上下文的推荐（同文件中的函数）
            file_context_query = """
            MATCH (current:Function), (recommended:Function)
            WHERE current.name CONTAINS $function_name 
               OR current.id CONTAINS $function_name
               AND recommended.file = current.file
               AND recommended <> current
            RETURN recommended.id as function_id,
                   recommended.name as function_name,
                   recommended.file as function_file,
                   1 as call_frequency,
                   'file_context' as recommendation_type,
                   '同文件中的其他函数' as reason
            LIMIT $limit
            """
            
            results = []
            
            # 执行所有查询并合并结果
            for query, query_name in [
                (direct_calls_query, "direct_calls"),
                (similar_functions_query, "similar_functions"), 
                (file_context_query, "file_context")
            ]:
                try:
                    query_results = session.run(query, 
                                              function_name=current_function,
                                              limit=limit)
                    for record in query_results:
                        results.append({
                            "function_id": record["function_id"],
                            "function_name": record["function_name"],
                            "function_file": record["function_file"],
                            "call_frequency": record["call_frequency"],
                            "recommendation_type": record["recommendation_type"],
                            "reason": record["reason"],
                            "query_source": query_name
                        })
                except Exception as e:
                    print(f"查询 {query_name} 执行失败: {e}")
            
            # 去重并按推荐度排序
            unique_results = {}
            for result in results:
                key = result["function_id"]
                if key not in unique_results or result["call_frequency"] > unique_results[key]["call_frequency"]:
                    unique_results[key] = result
            
            # 按调用频率排序
            sorted_results = sorted(unique_results.values(), 
                                  key=lambda x: x["call_frequency"], 
                                  reverse=True)
            
            return sorted_results[:limit]
    
    def get_next_function_calls(self, 
                              current_function: str,
                              limit: int = 5) -> List[Dict]:
        """
        获取当前函数可能调用的下一个函数
        
        Args:
            current_function: 当前函数名
            limit: 返回数量限制
            
        Returns:
            可能的下一个函数调用列表
        """
        with self.driver.session() as session:
            query = """
            MATCH (current:Function)-[c:CALLS]->(next:Function)
            WHERE current.name CONTAINS $function_name 
               OR current.id CONTAINS $function_name
            WITH next, count(c) as frequency, collect(c.line) as call_lines
            ORDER BY frequency DESC
            LIMIT $limit
            RETURN next.id as function_id,
                   next.name as function_name,
                   next.file as function_file,
                   frequency,
                   call_lines,
                   'next_call' as recommendation_type,
                   '当前函数经常调用的函数' as reason
            """
            
            results = []
            query_results = session.run(query, 
                                      function_name=current_function,
                                      limit=limit)
            
            for record in query_results:
                results.append({
                    "function_id": record["function_id"],
                    "function_name": record["function_name"],
                    "function_file": record["function_file"],
                    "frequency": record["frequency"],
                    "call_lines": record["call_lines"],
                    "recommendation_type": record["recommendation_type"],
                    "reason": record["reason"]
                })
            
            return results
    
    def get_function_signature_suggestions(self, 
                                         function_name: str,
                                         limit: int = 5) -> List[Dict]:
        """
        基于函数名获取可能的函数签名建议
        
        Args:
            function_name: 函数名（可能不完整）
            limit: 返回数量限制
            
        Returns:
            函数签名建议列表
        """
        with self.driver.session() as session:
            query = """
            MATCH (f:Function)
            WHERE f.name CONTAINS $function_name
               OR f.id CONTAINS $function_name
            WITH f, 
                 size([(f)-[:CALLS]->() | 1]) as outgoing_calls,
                 size([()-[:CALLS]->(f) | 1]) as incoming_calls
            ORDER BY (outgoing_calls + incoming_calls) DESC
            LIMIT $limit
            RETURN f.id as function_id,
                   f.name as function_name,
                   f.file as function_file,
                   outgoing_calls,
                   incoming_calls,
                   'signature_suggestion' as recommendation_type,
                   '基于函数名匹配的签名建议' as reason
            """
            
            results = []
            query_results = session.run(query, 
                                      function_name=function_name,
                                      limit=limit)
            
            for record in query_results:
                results.append({
                    "function_id": record["function_id"],
                    "function_name": record["function_name"],
                    "function_file": record["function_file"],
                    "outgoing_calls": record["outgoing_calls"],
                    "incoming_calls": record["incoming_calls"],
                    "recommendation_type": record["recommendation_type"],
                    "reason": record["reason"]
                })
            
            return results
    
    def get_context_aware_recommendations(self, 
                                        current_file: str,
                                        current_line: int,
                                        limit: int = 10) -> List[Dict]:
        """
        基于文件上下文和行号的智能推荐
        
        Args:
            current_file: 当前文件路径
            current_line: 当前行号
            limit: 返回数量限制
            
        Returns:
            上下文感知的推荐列表
        """
        with self.driver.session() as session:
            # 查询同文件中的函数调用关系
            query = """
            MATCH (caller:Function)-[c:CALLS]->(callee:Function)
            WHERE caller.file = $current_file
               OR callee.file = $current_file
               OR c.file = $current_file
            WITH caller, callee, c, 
                 abs(c.line - $current_line) as line_distance
            ORDER BY line_distance ASC, c.line ASC
            LIMIT $limit
            RETURN caller.id as caller_id,
                   caller.name as caller_name,
                   caller.file as caller_file,
                   callee.id as callee_id,
                   callee.name as callee_name,
                   callee.file as callee_file,
                   c.line as call_line,
                   line_distance,
                   'context_aware' as recommendation_type,
                   '基于文件位置和行号的上下文推荐' as reason
            """
            
            results = []
            query_results = session.run(query, 
                                      current_file=current_file,
                                      current_line=current_line,
                                      limit=limit)
            
            for record in query_results:
                results.append({
                    "caller_id": record["caller_id"],
                    "caller_name": record["caller_name"],
                    "caller_file": record["caller_file"],
                    "callee_id": record["callee_id"],
                    "callee_name": record["callee_name"],
                    "callee_file": record["callee_file"],
                    "call_line": record["call_line"],
                    "line_distance": record["line_distance"],
                    "recommendation_type": record["recommendation_type"],
                    "reason": record["reason"]
                })
            
            return results
    
    def get_popular_functions(self, 
                            file_pattern: str = None,
                            limit: int = 20) -> List[Dict]:
        """
        获取最受欢迎的函数（被调用次数最多）
        
        Args:
            file_pattern: 文件路径模式过滤
            limit: 返回数量限制
            
        Returns:
            受欢迎函数列表
        """
        with self.driver.session() as session:
            if file_pattern:
                query = """
                MATCH (f:Function)
                WHERE f.file CONTAINS $file_pattern
                WITH f, size([()-[:CALLS]->(f) | 1]) as incoming_calls
                ORDER BY incoming_calls DESC
                LIMIT $limit
                RETURN f.id as function_id,
                       f.name as function_name,
                       f.file as function_file,
                       incoming_calls,
                       'popular_function' as recommendation_type,
                       '最受欢迎的函数' as reason
                """
                query_results = session.run(query, 
                                          file_pattern=file_pattern,
                                          limit=limit)
            else:
                query = """
                MATCH (f:Function)
                WITH f, size([()-[:CALLS]->(f) | 1]) as incoming_calls
                ORDER BY incoming_calls DESC
                LIMIT $limit
                RETURN f.id as function_id,
                       f.name as function_name,
                       f.file as function_file,
                       incoming_calls,
                       'popular_function' as recommendation_type,
                       '最受欢迎的函数' as reason
                """
                query_results = session.run(query, limit=limit)
            
            results = []
            for record in query_results:
                results.append({
                    "function_id": record["function_id"],
                    "function_name": record["function_name"],
                    "function_file": record["function_file"],
                    "incoming_calls": record["incoming_calls"],
                    "recommendation_type": record["recommendation_type"],
                    "reason": record["reason"]
                })
            
            return results


def get_driver() -> Driver:
    """获取Neo4j驱动连接"""
    load_dotenv()
    uri = os.getenv("NEO4J_URI", "bolt://localhost:7687")
    user = os.getenv("NEO4J_USER", "neo4j")
    password = os.getenv("NEO4J_PASSWORD", "neo4j")
    return GraphDatabase.driver(uri, auth=(user, password))


def main():
    """演示代码推荐功能"""
    driver = get_driver()
    recommender = CodeRecommender(driver)
    
    try:
        print("=== 代码推荐系统演示 ===\n")
        
        # 演示1: 获取函数推荐
        print("1. 函数推荐示例:")
        recommendations = recommender.get_function_recommendations("main", limit=5)
        for i, rec in enumerate(recommendations, 1):
            print(f"  {i}. {rec['function_name']} ({rec['function_file']})")
            print(f"     理由: {rec['reason']}")
            print(f"     调用频率: {rec['call_frequency']}")
            print()
        
        # 演示2: 获取下一个函数调用
        print("2. 下一个函数调用建议:")
        next_calls = recommender.get_next_function_calls("main", limit=3)
        for i, call in enumerate(next_calls, 1):
            print(f"  {i}. {call['function_name']} (调用频率: {call['frequency']})")
            print(f"     调用行号: {call['call_lines']}")
            print()
        
        # 演示3: 获取受欢迎的函数
        print("3. 最受欢迎的函数:")
        popular = recommender.get_popular_functions(limit=5)
        for i, func in enumerate(popular, 1):
            print(f"  {i}. {func['function_name']} (被调用 {func['incoming_calls']} 次)")
            print(f"     文件: {func['function_file']}")
            print()
        
    finally:
        driver.close()


if __name__ == "__main__":
    main()
